import time
from io import BytesIO

import cv2
from PIL import Image
import onnxruntime
import numpy as np
import os
import torch
import base64
from ctypes import *


def get_files(input_path):
    file_list = []
    for root, dirs, filenames in os.walk(input_path):
        for filename in filenames:
            if filename.endswith((".jpg", ".tif", ".jpeg", ".png", ".JPG", ".TIF", ".JPEG", ".PNG")):
                file_list.append(root + "/" + filename)
    return file_list


class XueXin:
    def __init__(self, onnx_path="./new_weight2/xuexin_20230314_ok_simple.onnx"):
        self.alphaNum = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ+&~!@#'
        self.ort_session = onnxruntime.InferenceSession(onnx_path)

    def predict_img(self, img):
        img_np = np.divide(img, 255.0)
        img_np_pre = np.divide(np.subtract(img_np, 0.5), 0.5)
        img_np_pre = cv2.repeat(img_np_pre, 3, 1)
        img_np_pre = img_np_pre.reshape((1, 3, 100, 200))
        ort_inputs = {self.ort_session.get_inputs()[0].name: img_np_pre.astype(np.float32)}
        ort_outs = self.ort_session.run(None, ort_inputs)  # 50,1,69
        output = np.argmax(ort_outs[0].reshape(50, 69), axis=1)  # 50, 1
        output = output.reshape(1, 50)  # 1, 50
        predict = torch.unique_consecutive(torch.Tensor(output))
        predict = predict[predict != 68]
        label = ''.join([self.alphaNum[j] for j in predict.numpy().astype(np.int32).tolist()])
        return label

    def string_2_imgL(self):
        try:
            img_data = base64.b64decode(self.image_string)
            self.image = np.array(Image.open(BytesIO(img_data)).convert("L"))
            return True
        except Exception as e:
            self.image = str(e)
        return False


class XUEXIN_ONNX:
    def __init__(self, xuexin_onnx_path):
        self.MYDLL = CDLL("./ocr_server/libxuexin_onnx.so")

        #############################################
        #            input output arg
        #############################################
        self.MYDLL.XueXinInit.argtypes = [c_char_p]
        self.MYDLL.XueXinInit.restype = POINTER(c_ubyte)

        self.MYDLL.XueXinDetect.argtypes = [POINTER(c_ubyte), POINTER(c_ubyte), c_int, c_int, c_int, c_char_p]
        self.MYDLL.XueXinDetect.restype = c_void_p

        self.MYDLL.XueXinRelease.argtypes = [POINTER(c_ubyte)]
        self.MYDLL.XueXinRelease.restype = c_void_p
        self.xuexin_Handle = self.MYDLL.XueXinInit(c_char_p(xuexin_onnx_path))

    def predict_xuexin_onnx(self, image):
        if image.shape == (100, 200):
            height, width = image.shape
            src_pointer = image.ctypes.data_as(POINTER(c_ubyte))
            result = create_string_buffer(50)
            self.MYDLL.XueXinDetect(self.xuexin_Handle, src_pointer, height, width, width, result)
            pre_txt = result.value.decode()
            if '+' or '&' in pre_txt:
                pre_txt = str(eval(pre_txt[:-4].replace("&", "*")))
            return pre_txt
        return ""


if __name__ == '__main__':
    test_path = "E:/COMPANY/20200515__OCR_custom/ocr_server/verify_code_xuexin/train/root2_bin"
    file_list_ = get_files(test_path)
    xuexin_ = XueXin()
    for img_path in file_list_:
        image_gray = cv2.imread(img_path, 0)
        # print("image_gray=", image_gray)
        pre_y = xuexin_.predict_img(image_gray)

        random_str = str(time.time()).split('.')[0][:5]
        new_name = test_path+"_rename/" + pre_y + "_" + random_str+".png"
        try:
            os.rename(img_path, new_name)
        except Exception as e:
            print(e)
